#include "Drive.h"
#include <VersionHelpers.h>
#include <io.h>
#include <fstream>
#include "utils.h"
#include <Shlwapi.h>

#define DIR_FILE "C:\\Windows\\System32\\drivers\\netfilter2.sys"

std::string MakeMask(int s)
{
	IN_ADDR inaddr;
	inaddr.s_addr = 4294967295UL << (32 - s);
	char *txtIp = inet_ntoa(inaddr);
	std::string strip(txtIp);

	std::vector<std::string> is = std::split(strip, ".");
	std::string retmaskIP;
	if (is.size() == 4)
	{
		retmaskIP = is[3] + "." + is[2] + "." + is[1] + "." + is[0];
	}

	return retmaskIP;
}
bool NetDrive::DriveInit(bool grepLocaolIP)
{
	int ret = 0;

	m_NetApiModule = LoadLibraryA("nfapi.dll");
	if (m_NetApiModule <= 0)
	{
		m_NetApiModule = LoadLibraryA("./lib/nfapi.dll");
	}
	m_ProxyModule = LoadLibraryA("./lib/Proxy.dll");
	if (m_ProxyModule <= 0 || m_NetApiModule <= 0)
	{
		return false;
	}
	if (!FreeDrive())
	{
		return false;
	}
	typedef int (*Driver_Init)(char *, int, char *, char *);

	Driver_Init DriInit = (Driver_Init)GetProcAddress(m_ProxyModule, "Driver_Init");
	if (DriInit)
	{
		if (DriInit("0.0.0.0", 21999, "127.0.0.1", "[::1]") != 0)
		{
			return false;
		}
	}
	if (grepLocaolIP)
	{
		ret = ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "127.0.0.0", (CHAR *)MakeMask(8).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "10.0.0.0", (CHAR *)MakeMask(8).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "100.64.0.0", (CHAR *)MakeMask(10).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "169.254.0.0", (CHAR *)MakeMask(16).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "172.16.0.0", (CHAR *)MakeMask(12).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "192.0.0.0", (CHAR *)MakeMask(24).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "192.0.2.0", (CHAR *)MakeMask(24).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "192.88.99.0", (CHAR *)MakeMask(24).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "192.168.0.0", (CHAR *)MakeMask(16).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "198.18.0.0", (CHAR *)MakeMask(15).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "198.51.100.0", (CHAR *)MakeMask(24).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "203.0.113.0", (CHAR *)MakeMask(24).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "224.0.0.0", (CHAR *)MakeMask(4).c_str(), NF_ALLOW);
		ret |= ADDRule(false, 0, 0, 0, 0, 0, AF_INET, NULL, NULL, "240.0.0.0", (CHAR *)MakeMask(4).c_str(), NF_ALLOW);
	}

	ret |= ADDRule(0, 17, 0, 3, 0, 53, AF_INET, 0, 0, 0, 0, 0);  // DNS)
	ret |= ADDRule(0, 6, 0, 2, 0, 0, AF_INET, 0, 0, 0, 0, 16);//  ' TCP
	ret |= ADDRule(0, 17, 0, 0, 0, 0, AF_INET, 0, 0, 0, 0, 2);//  ' udp

	return ret;
}
int NetDrive::AddProxy(int type, int pid, char *ip, int port, char *user, char *passw, int mode)
{
	typedef int (*AddProxy)(int type, int pid, char *ip, int port, char *user, char *passw, int mode);
	AddProxy fun = (AddProxy)GetProcAddress(m_ProxyModule, "AddProxy");
	if (fun)
	{
		return fun(type, pid, ip, port, user, passw, mode);
	}
	return -1;
}
int NetDrive::DelteProxy(int pid)
{
	typedef int (*DelteProxy)(int pid);
	DelteProxy fun = (DelteProxy)GetProcAddress(m_ProxyModule, "DelteProxy");
	if (fun)
	{
		return fun(pid);
	}
	return -1;
}
std::string NetDrive::Formatting(long size)
{
	char buff[MAX_PATH] = {0};
	StrFormatByteSize64A(size, buff, MAX_PATH);
	return std::string(buff);
}
int NetDrive::DriverFree()
{
	typedef int (*Driver_Free)();
	Driver_Free fun = (Driver_Free)GetProcAddress(m_ProxyModule, "Driver_Free");
	if (fun)
	{
		fun();
	}
	if (m_NetApiModule)
	{
		FreeLibrary(m_NetApiModule);
	}
	if (m_ProxyModule)
	{
		FreeLibrary(m_ProxyModule);
	}
	return 0;
}

int NetDrive::SetDns(char *DnsAddr, int DnsPort, int OutTime)
{
	typedef int (*SetDns)(char *DnsAddr, int DnsPort, int OutTime);
	SetDns fun = (SetDns)GetProcAddress(m_ProxyModule, "SetDns");
	if (fun)
	{
		return fun(DnsAddr, DnsPort, OutTime);
	}
	return -1;
}
int NetDrive::SetDnsProxy(char *SSAddr, int SSPort, char *SSUser, char *SSPass)
{
	typedef int (*SetDnsProxy)(char *SSAddr, int SSPort, char *SSUser, char *SSPass);
	SetDnsProxy fun = (SetDnsProxy)GetProcAddress(m_ProxyModule, "SetDnsProxy");
	if (fun)
	{
		return fun(SSAddr, SSPort, SSUser, SSPass);
	}
	return -1;
}
long NetDrive::GetInfo(int pid, int type )
{
	typedef int (*GetInfo)(int pid, int type);
	GetInfo fun = (GetInfo)GetProcAddress(m_ProxyModule, "GetInfo");
	if (fun)
	{
		return fun(pid, type);
	}
	return -1;
}
int NetDrive::ADDRule(
	bool AddtoHead,
	INT protocol,
	INT processId,
	BYTE direction,
	SHORT localPort,
	SHORT remotePort,
	SHORT ip_family,
	CHAR *localIpAddress,
	CHAR *localIpAddressMask,
	CHAR *remoteIpAddress,
	CHAR *remoteIpAddressMask,
	INT filteringFlag)
{
	NF_RULE rule = {0};
	rule.protocol = protocol;
	rule.processId = processId;
	rule.direction = direction;
	rule.localPort = htons(localPort);
	rule.remotePort = htons(remotePort);
	rule.ip_family = ip_family;
	rule.filteringFlag = filteringFlag;
	if (localIpAddress)
	{
		memcpy(rule.localIpAddress, localIpAddress, IP_STR_LEM);
	}
	if (localIpAddress)
	{
		memcpy(rule.localIpAddressMask, localIpAddressMask, IP_STR_LEM);
	}
	if (localIpAddress)
	{
		memcpy(rule.remoteIpAddress, remoteIpAddress, IP_STR_LEM);
	}
	if (localIpAddress)
	{
		memcpy(rule.remoteIpAddressMask, remoteIpAddressMask, IP_STR_LEM);
	}
	typedef int (*nf_addRule)(NF_RULE*, bool);
	nf_addRule addRule = (nf_addRule)GetProcAddress(m_NetApiModule, "nf_addRule");
	if (addRule)
	{
		return addRule(&rule, AddtoHead);
	}
	return -1;
}
bool NetDrive::FreeDrive()
{

	SYSTEM_INFO si;
	bool is64bits;
	PVOID OffOldValue;
	GetNativeSystemInfo(&si);

	if (si.wProcessorArchitecture == PROCESSOR_ARCHITECTURE_AMD64 ||
		si.wProcessorArchitecture != PROCESSOR_ARCHITECTURE_IA64)
	{
		is64bits = true; //64
	}
	else
	{
		is64bits = false; // 32
	}

	if (Wow64DisableWow64FsRedirection(&OffOldValue))
	{
		if (IsWindows10OrGreater() ||
			IsWindows8OrGreater() ||
			IsWindowsServer())
		{
			if ((_access(DIR_FILE, 0)) == -1)
			{
				if (is64bits)
				{
					CopyFile(DIR_FILE, "./lib/natsec_Win10x64.sys", FALSE);
				}
				else
				{
					CopyFile(DIR_FILE, "./lib/natsec_Win10x86.sys", FALSE);
				}
			}
		}
		else if (IsWindows7OrGreater())
		{
			if ((_access(DIR_FILE, 0)) == -1)
			{
				if (is64bits)
				{
					CopyFile(DIR_FILE, "./lib/natsec_Win7x64.sys", FALSE);
				}
				else
				{
					CopyFile(DIR_FILE, "./lib/natsec_Win7x86.sys", FALSE);
				}
			}
		}
		else
		{
			if ((_access(DIR_FILE, 0)) == -1)
			{
				if (is64bits)
				{
					CopyFile(DIR_FILE, "./lib/natsec_Xpx64.sys", FALSE);
				}
				else
				{
					CopyFile(DIR_FILE, "./lib/natsec_Xpx86.sys", FALSE);
				}
			}
		}
		if (OffOldValue == NULL)
		{
			if (!Wow64RevertWow64FsRedirection(&OffOldValue))
			{
				return false;
			}
		}
		return true;
	}
	return false;
}
NetDrive::~NetDrive()
{
	DriverFree();
}